import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import sys
import json
import os

# 定义与训练时一致的模型架构
class TumorClassifier(nn.Module):
    def __init__(self, num_classes):
        super(TumorClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(32 * 56 * 56, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.classifier(x)
        return x

# 检查命令行参数
if len(sys.argv) != 2:
    print(json.dumps({"error": "请提供图像文件路径"}))
    sys.exit(1)

image_path = sys.argv[1]

# 检查文件是否存在
if not os.path.exists(image_path):
    print(json.dumps({"error": "图像文件不存在"}))
    sys.exit(1)

# 检查文件是否为图像格式
try:
    img = Image.open(image_path)
except Exception as e:
    print(json.dumps({"error": "无法打开图像文件"}))
    sys.exit(1)

# 定义数据预处理步骤（与训练时保持一致）
def get_transforms():
    return transforms.Compose([
        transforms.Resize((224, 224)),  # 调整图像大小
        transforms.ToTensor(),  # 转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化
    ])

# 类别标签
classes = ['glioma', 'meningioma', 'notumor', 'pituitary']

# 创建模型并加载权重
def load_model():
    # 创建与训练时一致的模型架构
    model = TumorClassifier(num_classes=4)
    
    # 加载训练好的权重
    try:
        model.load_state_dict(torch.load('best_model.pth', map_location=torch.device('cpu')))
    except Exception as e:
        print(json.dumps({"error": f"无法加载模型权重: {str(e)}"}))
        sys.exit(1)
    
    model.eval()  # 设置为评估模式
    return model

# 预测函数
def predict(image_path):
    # 加载模型
    model = load_model()
    
    # 加载并预处理图像
    image = Image.open(image_path).convert('RGB')
    transform = get_transforms()
    image_tensor = transform(image).unsqueeze(0)  # 添加批次维度
    
    # 进行预测
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = torch.nn.functional.softmax(outputs, dim=1)
        confidence, predicted = torch.max(probabilities, 1)
    
    # 获取预测结果
    predicted_class = classes[predicted.item()]
    confidence_score = confidence.item() * 100  # 转换为百分比
    
    # 获取所有类别的概率
    all_probabilities = []
    for i, prob in enumerate(probabilities[0]):
        all_probabilities.append({
            "class": classes[i],
            "confidence": prob.item() * 100
        })
    
    # 按置信度排序
    all_probabilities.sort(key=lambda x: x["confidence"], reverse=True)
    
    # 返回结果
    return {
        "predicted_class": predicted_class,
        "confidence": confidence_score,
        "all_probabilities": all_probabilities
    }

# 执行预测并输出结果
try:
    result = predict(image_path)
    print(json.dumps(result))
except Exception as e:
    print(json.dumps({"error": f"预测过程中发生错误: {str(e)}"}))
    sys.exit(1)